import time
import urllib.error
import urllib.parse
import urllib.request
import ffmpeg_streaming
import requests
from ffmpeg_streaming import Formats

import amemv

cookie: str = ''
# Don't need to write your cookie. A static cookie will skip the
# cold-boot stage.
web_rid: str = '1962143378'
savepath: str = ''
svc_key: str = '' # Fill your ServerChan sendkey here.


def PushMsg(type: str, web_rid: str, cookie: str, renew: bool):
    rid = amemv.RequestRoomID(web_rid=web_rid, cookie=cookie)
    nickname = amemv.RequestInformation(id="nickname", room_id=rid)
    signature = amemv.RequestInformation(id="signature", room_id=rid)
    signature = signature.replace("\\n", "\n\n")
    title_web = amemv.RequestInformation(id="title", room_id=rid)
    svcurl = "https://sctapi.ftqq.com/" + svc_key + '.send'
    if type == 'StartListen':
        type = '正在监听'
    if type == 'StartRec':
        type = '开始录制'
    if type == 'Finish':
        type = '录制结束'
    if renew:
        cookie_str: str = "Cookie 将在下一次监听之前自动刷新，以避免抖音风控。"
    else:
        cookie_str: str = "Cookie 为固定值，建议您留空允许脚本自动更新，或建议您下播后手动更新。"
    title = '[amemv] 抖音直播间“' + nickname + '” ' + type + '捏~'
    message = f"""直播间标题：{title_web}

直播间地址：https://live.douyin.com/{web_rid}

直播者签名：

{signature}

---

本次录制使用凭据：

```{cookie}```

凭据提示：{cookie_str}"""
    requests.post(url=svcurl, params={'desp': message, 'title': title})


def StreamStatus(live_url):
    try:
        urllib.request.urlopen(live_url)
    except urllib.error.HTTPError:
        return False
    else:
        return True


def roomIdStatus(web_rid):
    before = amemv.RequestRoomID(web_rid=web_rid, cookie=cookie)
    time.sleep(5)
    after = amemv.RequestRoomID(web_rid=web_rid, cookie=cookie)
    if before == after:
        return False
    else:
        return True


def RecordStream(live_url: str, nickname: str):
    timestr = time.strftime("%Y-%m-%d--%H_%M_%S", time.localtime())
    savestr = savepath + "//" + "抖音录制_" + nickname + "_" + timestr + ".flv"
    stream = ffmpeg_streaming.input(live_url)

    video = stream.stream2file(Formats.hevc())
    video.output(savestr)


if __name__ == '__main__':
    predef: bool = False
    renew: bool = True
    if cookie == '':
        print("[!] amemv-rec starting from cold-and-dark")
        print("[#] Generating new cookies...")
        cookie = amemv.RenewCookies()
        bootstrap = True
    else:
        print("[!] amemv-rec starting with a exist cookie token.")
        predef = True
        bootstrap = True
    print("[i] Cookies = " + cookie)
    rid = amemv.RequestRoomID(web_rid=web_rid, cookie=cookie)
    nickname = amemv.RequestInformation(id="nickname", room_id=rid)
    print("[i] Target room nickname = " + nickname)
    while True:
        status = bool(True)
        if predef == True:
            renew = False

        if bootstrap == False and renew == True:
            cookie = amemv.RenewCookies()

        # Check
        while True:
            # Check availability of stream.
            # If true, start recording immediately
            roomId = amemv.RequestRoomID(web_rid=web_rid, cookie=cookie)
            url = amemv.RequestInformation(id="rtmp_pull_url", room_id=roomId)
            streamUrl = amemv.RequestInformation(id="rtmp_pull_url", room_id=roomId)
            if StreamStatus(streamUrl):
                print("[!] Target room already on live.")
                print("[#] Starting recording.")
                PushMsg(type="StartRec", web_rid=web_rid, cookie=cookie,renew=renew)
                RecordStream(live_url=url, nickname=nickname)
                PushMsg(type="Finish", web_rid=web_rid, cookie=cookie,renew=renew)
                break
            else:
                print("[!] Target room NOT on live.")
                print("[#] Waiting for live start.")
                PushMsg(type="StartListen", web_rid=web_rid, cookie=cookie,renew=renew)
                break

            # If stream can't catch, waiting for new roomId
        while True:

            if roomIdStatus(web_rid=web_rid):
                print("[!] Target room on live.")
                print("[#] Starting recording.")
                PushMsg(type="StartRec", web_rid=web_rid, cookie=cookie,renew=renew)
                RecordStream(live_url=amemv.RequestInformation(id="rtmp_pull_url",
                                                               room_id=amemv.RequestRoomID(web_rid=web_rid,
                                                                                           cookie=cookie)),
                             nickname=nickname)
                PushMsg(type="Finish", web_rid=web_rid, cookie=cookie, renew=renew)
                break
            else:
                continue
